#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
测试聊天功能（支持打字和语音输入）
"""

import asyncio
import sys
import threading
import queue
from pathlib import Path
import pygame
import io
import os
import atexit
import logging
from dotenv import load_dotenv
from collections import deque

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 注册程序退出时的清理函数
# 全局音频播放队列和控制变量
audio_queue = deque()
audio_playing = False
audio_thread = None
audio_lock = threading.Lock()

# 打断事件与共享变量（用于跨线程通知与传递文本）# 全局变量：打断事件
barge_in_event = threading.Event()
barge_in_text = None
barge_in_lock = threading.Lock()

# 全局变量：AI生成状态
ai_generating = False
ai_generating_lock = threading.Lock()

def cleanup_audio():
    """清理音频资源"""
    global audio_playing, audio_thread
    audio_playing = False
    if audio_thread and audio_thread.is_alive():
        audio_thread.join(timeout=1.0)
    if pygame.mixer.get_init():
        pygame.mixer.quit()
        
atexit.register(cleanup_audio)

# 加载环境变量
load_dotenv()

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

from src.clients.chat_client import LLMChatClient, ChatConfig
from src.clients.tts_manager import TTSManager
from src.clients.voice_to_text import RealtimeAsrClient


async def interrupt_tts_playback(chat_client=None, tts_manager=None):
    """立即停止TTS播放并清空队列（用于打断），支持取消TTS合成和LLM流"""
    try:
        # 设置打断事件
        barge_in_event.set()
        
        # 取消正在进行的TTS合成
        if tts_manager:
            try:
                await tts_manager.cancel_synthesis()
                print("🛑 TTS合成已取消")
            except Exception as e:
                print(f"⚠️ TTS取消失败: {e}")
        
        # 取消正在进行的LLM流式生成
        if chat_client:
            try:
                await chat_client.cancel_stream()
                print("🛑 LLM流式生成已取消")
            except Exception as e:
                print(f"⚠️ LLM流取消失败: {e}")
        
        # 停止当前音频播放并清空队列
        with audio_lock:
            audio_queue.clear()
        if pygame.mixer.get_init():
            pygame.mixer.stop()
        print("⏹️ 已停止当前TTS播放并清空队列")
    except Exception as e:
        print(f"❌ 停止音频播放失败: {e}")


def _audio_player_thread():
    """音频播放线程，按队列顺序播放"""
    global audio_playing, audio_queue
    
    while audio_playing:
        try:
            with audio_lock:
                if not audio_queue:
                    continue
                audio_data = audio_queue.popleft()
            
            # 检查mixer是否已初始化
            if not pygame.mixer.get_init():
                pygame.mixer.pre_init(frequency=44100, size=-16, channels=2, buffer=256)
                pygame.mixer.init()
            
            # 播放音频
            audio_io = io.BytesIO(audio_data)
            sound = pygame.mixer.Sound(audio_io)
            
            print("🔊 开始播放AI回复...")
            sound.play()
            
            # 等待当前音频播放完成
            while pygame.mixer.get_busy():
                pygame.time.wait(10)
                if not audio_playing:  # 支持中断（程序退出）
                    break
                # 如果发生打断事件，立即停止播放
                if barge_in_event.is_set():
                    pygame.mixer.stop()
                    break
            
            print("✅ 音频播放完成")
            
        except Exception as e:
            print(f"❌ 音频播放失败: {e}")
        
        # 短暂休眠避免CPU占用过高
        pygame.time.wait(10)


def play_audio_data(audio_data: bytes):
    """播放音频数据（队列播放方式）"""
    global audio_playing, audio_thread, audio_queue
    
    try:
        # 将音频数据加入队列
        with audio_lock:
            audio_queue.append(audio_data)
        
        # 启动播放线程（如果尚未启动）
        if not audio_playing:
            audio_playing = True
            audio_thread = threading.Thread(target=_audio_player_thread, daemon=True)
            audio_thread.start()
            print("🎵 音频播放线程已启动")
        
        print(f"📝 音频已加入播放队列，当前队列长度: {len(audio_queue)}")
        
    except Exception as e:
        print(f"❌ 音频队列添加失败: {e}")


async def process_ai_response(chat_client, tts_manager, user_message):
    """处理AI回复并合成语音（实时打印+实时TTS，支持打断）"""
    global ai_generating
    
    print(f"\n👤 用户: {user_message}")
    print("🤖 AI正在思考...")
    print("🤖 AI回复: ", end="", flush=True)  # 开始实时打印
    
    # 设置AI生成状态
    with ai_generating_lock:
        ai_generating = True
    
    # 每轮开始清除上一次可能残留的打断事件
    barge_in_event.clear()
    
    # 重置TTS和聊天客户端的取消标志
    try:
        if hasattr(tts_manager, 'reset_cancel_flag'):
            tts_manager.reset_cancel_flag()
        if hasattr(chat_client, 'reset_cancel_flag'):
            chat_client.reset_cancel_flag()
    except Exception as e:
        print(f"⚠️ 重置取消标志失败: {e}")
    
    # 配置参数
    MAX_SENTENCE_LENGTH = 80  # 最大句子长度（适配TTS性能）
    MIN_SENTENCE_LENGTH = 8   # 最小句子长度（防止过短片段）
    PRIMARY_DELIMITERS = ['。', '！', '？', '.', '!', '?']  # 主要句子结束符
    SECONDARY_DELIMITERS = ['，', ',', '；', ';', '：', ':', '、', '\n']  # 次级分隔符
    DEBOUNCE_DELAY = 0.3  # 防抖延迟（秒）
    
    def clean_text_for_tts(text: str) -> str:
        """清理文本，移除TTS不支持的字符"""
        import re
        # 移除表情符号和特殊字符，保留中文、英文、数字和基本标点
        cleaned = re.sub(r'[^\u4e00-\u9fff\u3400-\u4dbf\w\s，。！？,.!?；;：:（）()【】\[\]""''"\'-]', '', text)
        return cleaned.strip()

    def is_sentence_complete(text: str) -> bool:
        """判断文本是否构成完整句子"""
        text = text.strip()
        if not text:
            return False
        
        # 检查是否以句子结束符结尾
        if any(text.endswith(delimiter) for delimiter in PRIMARY_DELIMITERS):
            return len(text) >= MIN_SENTENCE_LENGTH
        
        return False
    
    def split_long_sentence(text: str, max_length: int) -> list:
        """智能分割过长的句子"""
        if len(text) <= max_length:
            return [text]
        
        # 优先按次级分隔符分割
        for delimiter in SECONDARY_DELIMITERS:
            if delimiter in text:
                parts = text.split(delimiter)
                result = []
                current_part = ""
                
                for i, part in enumerate(parts):
                    part = part.strip()
                    if not part:
                        continue
                        
                    test_part = current_part + (delimiter if current_part else "") + part
                    
                    if len(test_part) <= max_length:
                        current_part = test_part
                    else:
                        if current_part and len(current_part) >= MIN_SENTENCE_LENGTH:
                            result.append(current_part)
                        current_part = part
                
                if current_part and len(current_part) >= MIN_SENTENCE_LENGTH:
                    result.append(current_part)
                
                if len(result) > 1:  # 成功分割
                    return result
        
        # 如果无法按分隔符分割，按字符数智能分割（尽量在词边界）
        result = []
        current_pos = 0
        
        while current_pos < len(text):
            end_pos = min(current_pos + max_length, len(text))
            
            # 如果不是最后一段，尝试在词边界分割
            if end_pos < len(text):
                # 向前查找合适的分割点
                for i in range(end_pos - 1, current_pos + MIN_SENTENCE_LENGTH - 1, -1):
                    if text[i] in [' ', '\n', '，', ',', '、']:
                        end_pos = i + 1
                        break
            
            segment = text[current_pos:end_pos].strip()
            if segment and len(segment) >= MIN_SENTENCE_LENGTH:
                result.append(segment)
            
            current_pos = end_pos
        
        return result if result else [text]
    
    try:
        # 获取AI流式回复
        full_response = ""
        sentence_buffer = ""
        interrupted = False
        last_sentence_time = 0  # 防抖时间戳
        
        async for chunk in chat_client.send_message_stream(user_message):
            # 检测是否发生打断，若已打断则停止后续生成
            if barge_in_event.is_set():
                interrupted = True
                await interrupt_tts_playback(chat_client, tts_manager)
                print("\n⛔ 已检测到用户打断，停止本轮回复。")
                break
            
            if chunk:
                # 实时打印AI回复字符
                print(chunk, end="", flush=True)
                
                full_response += chunk
                sentence_buffer += chunk
                
                # 检查是否有完整的句子（增加防抖机制）
                current_time = asyncio.get_event_loop().time()
                if any(punct in sentence_buffer for punct in PRIMARY_DELIMITERS):
                    # 找到句子结束位置
                    for i, char in enumerate(sentence_buffer):
                        if char in PRIMARY_DELIMITERS:
                            # 若出现打断，终止处理
                            if barge_in_event.is_set():
                                interrupted = True
                                await interrupt_tts_playback(chat_client, tts_manager)
                                print("\n⛔ 已检测到用户打断，停止本轮回复。")
                                break
                            
                            sentence = sentence_buffer[:i+1].strip()
                            
                            # 防抖检查：避免过于频繁的TTS调用
                            if sentence and is_sentence_complete(sentence):
                                if current_time - last_sentence_time >= DEBOUNCE_DELAY:
                                    last_sentence_time = current_time
                                    
                                    # 检查句子长度，如果过长则分割
                                    if len(sentence) > MAX_SENTENCE_LENGTH:
                                        parts = split_long_sentence(sentence, MAX_SENTENCE_LENGTH)
                                        for part in parts:
                                            if barge_in_event.is_set():
                                                interrupted = True
                                                await interrupt_tts_playback(chat_client, tts_manager)
                                                print("\n⛔ 已检测到用户打断，停止本轮回复。")
                                                break
                                            if part.strip():
                                                cleaned_part = clean_text_for_tts(part.strip())
                                                if cleaned_part and len(cleaned_part) >= MIN_SENTENCE_LENGTH:
                                                    print(f"🎯 合成分段: {cleaned_part}")
                                                    try:
                                                        audio_data = await tts_manager.synthesize_async(cleaned_part)
                                                        # 合成完成后再次检查是否打断，若已打断则不入队
                                                        if barge_in_event.is_set():
                                                            interrupted = True
                                                            await interrupt_tts_playback(chat_client, tts_manager)
                                                            print("\n⛔ 已检测到用户打断，停止本轮回复。")
                                                            break
                                                        play_audio_data(audio_data)
                                                    except Exception as e:
                                                        print(f"❌ 语音合成失败: {e}")
                                        if interrupted:
                                            break
                                    else:
                                        cleaned_sentence = clean_text_for_tts(sentence)
                                        if cleaned_sentence and len(cleaned_sentence) >= MIN_SENTENCE_LENGTH:
                                            print(f"🎯 合成句子: {cleaned_sentence}")
                                            try:
                                                audio_data = await tts_manager.synthesize_async(cleaned_sentence)
                                                if barge_in_event.is_set():
                                                    interrupted = True
                                                    await interrupt_tts_playback(chat_client, tts_manager)
                                                    print("\n⛔ 已检测到用户打断，停止本轮回复。")
                                                    break
                                                play_audio_data(audio_data)
                                            except Exception as e:
                                                print(f"❌ 语音合成失败: {e}")
                                
                                # 移除已处理的句子
                                sentence_buffer = sentence_buffer[i+1:]
                                break
                    if interrupted:
                        break
        
        # 若未打断，处理剩余的文本
        if not interrupted and sentence_buffer.strip():
            remaining_text = sentence_buffer.strip()
            # 只处理有意义长度的剩余文本
            if len(remaining_text) >= MIN_SENTENCE_LENGTH:
                if len(remaining_text) > MAX_SENTENCE_LENGTH:
                    parts = split_long_sentence(remaining_text, MAX_SENTENCE_LENGTH)
                    for part in parts:
                        if barge_in_event.is_set():
                            interrupted = True
                            await interrupt_tts_playback(chat_client, tts_manager)
                            print("\n⛔ 已检测到用户打断，停止本轮回复。")
                            break
                        if part.strip():
                            cleaned_part = clean_text_for_tts(part.strip())
                            if cleaned_part and len(cleaned_part) >= MIN_SENTENCE_LENGTH:
                                print(f"🎯 合成剩余分段: {cleaned_part}")
                                try:
                                    audio_data = await tts_manager.synthesize_async(cleaned_part)
                                    if barge_in_event.is_set():
                                        interrupted = True
                                        await interrupt_tts_playback(chat_client, tts_manager)
                                        print("\n⛔ 已检测到用户打断，停止本轮回复。")
                                        break
                                    play_audio_data(audio_data)
                                except Exception as e:
                                    print(f"❌ 语音合成失败: {e}")
                else:
                    cleaned_remaining = clean_text_for_tts(remaining_text)
                    if cleaned_remaining and len(cleaned_remaining) >= MIN_SENTENCE_LENGTH:
                        print(f"🎯 合成剩余文本: {cleaned_remaining}")
                        try:
                            audio_data = await tts_manager.synthesize_async(cleaned_remaining)
                            if not barge_in_event.is_set():
                                play_audio_data(audio_data)
                        except Exception as e:
                            print(f"❌ 语音合成失败: {e}")
            else:
                cleaned_remaining = clean_text_for_tts(remaining_text)
                if cleaned_remaining and not barge_in_event.is_set():
                    print(f"🎯 合成剩余文本: {cleaned_remaining}")
                    try:
                        audio_data = await tts_manager.synthesize_async(cleaned_remaining)
                        if not barge_in_event.is_set():
                            play_audio_data(audio_data)
                        else:
                            interrupted = True
                            await interrupt_tts_playback(chat_client, tts_manager)
                            print("\n⛔ 已检测到用户打断，停止本轮回复。")
                    except Exception as e:
                        print(f"❌ 语音合成失败: {e}")
        
        if not interrupted:
            print(f"\n\n📝 AI回复完成，共 {len(full_response)} 字符")
            # 记录AI回复内容到日志
            logger.info(full_response)
        else:
            print("\n🛑 本轮AI回复已因用户打断而提前结束。")
        
        return full_response
        
    except Exception as e:
        print(f"❌ AI回复处理失败: {e}")
        import traceback
        traceback.print_exc()
        return None
    finally:
        # 清除AI生成状态
        with ai_generating_lock:
            ai_generating = False

class VoiceInputHandler:
    """语音输入处理器"""
    
    def __init__(self):
        self.voice_queue = queue.Queue()
        self.asr_client = None
        self.is_listening = False
        
    def on_voice_result(self, text: str, is_final: bool):
        """语音识别结果回调"""
        # 修改逻辑：当文本长度足够时就处理，不必等待is_final=True
        if text.strip() and len(text.strip()) >= 2:  # 至少2个字符才处理
            print(f"\n🎤 语音识别: {text} (final={is_final})")
            # 避免重复添加相同的文本
            if not hasattr(self, '_last_text') or self._last_text != text.strip():
                self.voice_queue.put(text.strip())
                self._last_text = text.strip()
                print(f"✅ 语音输入已添加到队列: {text.strip()}")
            
            # 扩展打断逻辑：不仅在播放时打断，在AI思考/生成时也允许打断
            try:
                # 检查是否需要触发打断
                should_interrupt = False
                interrupt_reason = ""
                
                # 情况1：正在播放音频或有待播音频
                if pygame.mixer.get_busy() or len(audio_queue) > 0:
                    should_interrupt = True
                    interrupt_reason = "音频播放中"
                
                # 情况2：检查是否有AI正在生成回复
                elif ai_generating:
                    should_interrupt = True
                    interrupt_reason = "AI生成中"
                
                if should_interrupt:
                    global barge_in_text
                    with barge_in_lock:
                        barge_in_text = text.strip()
                    if not barge_in_event.is_set():
                        print(f"⛔ 检测到用户打断请求（{interrupt_reason}），正在停止AI生成和TTS播放...")
                    barge_in_event.set()
            except Exception:
                pass
        
    def on_voice_error(self, error: str):
        """语音识别错误回调"""
        print(f"❌ 语音识别错误: {error}")
        
    async def start_voice_recognition(self):
        """启动语音识别"""
        if self.is_listening:
            return
            
        try:
            # 从环境变量获取WebSocket URL
            voice_url = os.getenv('VOICE_WEBSOCKET_URL')
            if not voice_url:
                raise ValueError("VOICE_WEBSOCKET_URL 环境变量未设置")
                
            print("🎤 语音识别已启动，请说话...")
            
            # 使用 async context manager 正确初始化语音识别客户端
            async with RealtimeAsrClient(
                url=voice_url,
                on_result=self.on_voice_result,
                on_error=self.on_voice_error
            ) as asr_client:
                self.asr_client = asr_client
                self.is_listening = True
                await asr_client.start()
        except Exception as e:
            print(f"❌ 启动语音识别失败: {e}")
            self.is_listening = False
            
    async def stop_voice_recognition(self):
        """停止语音识别"""
        if self.asr_client and self.is_listening:
            await self.asr_client.stop()
            self.is_listening = False
            print("🎤 语音识别已停止")
            
    def get_voice_input(self):
        """获取语音输入"""
        try:
            return self.voice_queue.get_nowait()
        except queue.Empty:
            return None

async def get_user_input_mode():
    """获取用户输入模式选择"""
    print("\n📝 请选择输入模式:")
    print("1. 打字输入")
    print("2. 语音输入")
    print("3. 混合模式（打字+语音）")
    
    while True:
        try:
            choice = input("请输入选择 (1/2/3): ").strip()
            if choice in ['1', '2', '3']:
                return int(choice)
            else:
                print("❌ 请输入有效选择 (1/2/3)")
        except KeyboardInterrupt:
            return None

async def handle_text_input():
    """处理文字输入"""
    try:
        user_input = input("\n👤 请输入消息: ").strip()
        return user_input if user_input else None
    except EOFError:
        return None

async def main():
    """主函数"""
    print("💬 智能聊天程序（支持打字和语音输入）")
    print("=" * 60)
    print("功能: 支持文字输入和语音输入，AI会回复并播放语音")
    print("命令: 输入 'quit'、'exit' 或 '退出' 结束程序")
    print("     输入 'switch' 或 '切换' 切换输入模式")
    print("=" * 60)
    
    try:
        # 获取输入模式
        input_mode = await get_user_input_mode()
        if input_mode is None:
            print("\n👋 程序被用户中断，再见！")
            return
            
        mode_names = {1: "打字输入", 2: "语音输入", 3: "混合模式"}
        print(f"\n✅ 已选择: {mode_names[input_mode]}")
        
        # 初始化TTS管理器
        print("🔧 初始化TTS管理器...")
        tts_manager = TTSManager()
        
        print("🔧 初始化聊天客户端...")
        config = ChatConfig()
        
        # 初始化语音输入处理器
        voice_handler = None
        if input_mode in [2, 3]:
            voice_handler = VoiceInputHandler()
        
        # 使用async with语句管理聊天客户端
        async with LLMChatClient(config) as chat_client:
            print("✅ 初始化完成！")
            
            # 启动语音识别（如果需要）
            voice_task = None
            if voice_handler:
                try:
                    print("🎤 启动语音识别...")
                    voice_task = asyncio.create_task(voice_handler.start_voice_recognition())
                    await asyncio.sleep(1)  # 等待语音识别启动
                except Exception as e:
                    print(f"❌ 语音识别启动失败: {e}")
                    if input_mode == 2:  # 纯语音模式失败则退出
                        return
                    input_mode = 1  # 降级到打字模式
                    voice_handler = None
            
            print("\n🎯 开始聊天！")
            if input_mode == 1:
                print("💡 提示: 直接输入文字消息")
            elif input_mode == 2:
                print("💡 提示: 对着麦克风说话")
            else:
                print("💡 提示: 可以输入文字或对着麦克风说话")
            
            # 聊天循环
            while True:
                try:
                    user_input = None
                    
                    # 根据模式获取用户输入
                    if input_mode == 1:  # 纯打字模式
                        user_input = await handle_text_input()
                        
                    elif input_mode == 2:  # 纯语音模式
                        if voice_handler:
                            print("\n🎤 请说话...")
                            # 等待语音输入
                            while True:
                                voice_input = voice_handler.get_voice_input()
                                if voice_input:
                                    user_input = voice_input
                                    break
                                await asyncio.sleep(0.1)
                                
                    else:  # 混合模式
                        print("\n📝 请输入消息（文字）或直接说话（语音）:")
                        print("💡 提示: 按回车键输入文字，或直接对麦克风说话")
                        
                        # 创建文字输入任务
                        async def get_text_input():
                            loop = asyncio.get_event_loop()
                            return await loop.run_in_executor(None, input, "👤 文字输入: ")
                        
                        # 同时等待文字和语音输入
                        text_task = asyncio.create_task(get_text_input())
                        
                        try:
                            while True:
                                # 检查语音输入
                                if voice_handler:
                                    voice_input = voice_handler.get_voice_input()
                                    if voice_input:
                                        user_input = voice_input
                                        text_task.cancel()  # 取消文字输入任务
                                        print(f"🎤 语音识别: {voice_input}")
                                        break
                                
                                # 检查文字输入任务是否完成
                                if text_task.done():
                                    try:
                                        text_input = text_task.result().strip()
                                        if text_input:
                                            user_input = text_input
                                            break
                                    except asyncio.CancelledError:
                                        pass
                                    except Exception:
                                        pass
                                
                                await asyncio.sleep(0.1)
                        except Exception as e:
                            text_task.cancel()
                            print(f"❌ 输入处理出错: {e}")
                            continue
                    
                    # 处理用户输入
                    if not user_input:
                        continue
                        
                    # 检查退出命令
                    if user_input.lower() in ['quit', 'exit', '退出']:
                        print("👋 再见！")
                        break
                        
                    # 检查切换模式命令
                    if user_input.lower() in ['switch', '切换']:
                        print("\n🔄 切换输入模式...")
                        # 停止当前语音识别
                        if voice_handler:
                            await voice_handler.stop_voice_recognition()
                            if voice_task:
                                voice_task.cancel()
                        
                        # 重新选择模式
                        new_mode = await get_user_input_mode()
                        if new_mode and new_mode != input_mode:
                            input_mode = new_mode
                            print(f"✅ 已切换到: {mode_names[input_mode]}")
                            
                            # 重新初始化语音处理器
                            if input_mode in [2, 3]:
                                voice_handler = VoiceInputHandler()
                                voice_task = asyncio.create_task(voice_handler.start_voice_recognition())
                                await asyncio.sleep(1)
                            else:
                                voice_handler = None
                        continue
                    
                    # 处理AI回复
                    await process_ai_response(chat_client, tts_manager, user_input)
                    
                    # 若发生打断，则立即处理最新的语音输入，避免等待下一轮循环
                    if 'barge_in_event' in globals() and barge_in_event.is_set():
                        try:
                            new_text = None
                            if 'barge_in_lock' in globals() and 'barge_in_text' in globals():
                                with barge_in_lock:
                                    new_text = barge_in_text
                                    # 清空已消费的打断文本
                                    globals()['barge_in_text'] = None
                            # 重置事件，准备下一轮
                            barge_in_event.clear()
                            if new_text and new_text.strip():
                                print(f"\n➡️ 检测到打断后的新语音输入：{new_text}")
                                await process_ai_response(chat_client, tts_manager, new_text.strip())
                                continue
                        except Exception:
                            pass
                except KeyboardInterrupt:
                    print("\n\n👋 程序被用户中断，再见！")
                    break
                except Exception as e:
                    print(f"❌ 处理消息时出错: {e}")
                    continue
            
            # 清理语音识别
            if voice_handler:
                await voice_handler.stop_voice_recognition()
            if voice_task:
                voice_task.cancel()
                
    except Exception as e:
        print(f"❌ 程序初始化失败: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    asyncio.run(main())